import cv2
import numpy as np

# 定义要检测的颜色及其HSV范围
color_config = {
    "red": {"lower": [0, 100, 100], "upper": [10, 255, 255]},  # 红色（注意处理色相环绕）
    "black": {"lower": [0, 0, 0], "upper": [180, 80, 190]},  # 黑色
    "blue": {"lower": [90, 136, 20], "upper": [115, 255, 255]},  # 蓝色
    "yellow": {"lower": [17, 73, 148], "upper": [73, 181, 255]}  # 黄色
}


def detect_colored_polygons(image):
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    result = image.copy()

    for color_name, config in color_config.items():
        # 处理红色（色相在0-10和170-180）
        if color_name == "red":
            lower1 = np.array(config["lower"], dtype=np.uint8)
            upper1 = np.array(config["upper"], dtype=np.uint8)
            lower2 = np.array([170, config["lower"][1], config["lower"][2]], dtype=np.uint8)
            upper2 = np.array([180, config["upper"][1], config["upper"][2]], dtype=np.uint8)
            mask1 = cv2.inRange(hsv, lower1, upper1)
            mask2 = cv2.inRange(hsv, lower2, upper2)
            mask = cv2.bitwise_or(mask1, mask2)
        else:
            lower = np.array(config["lower"], dtype=np.uint8)
            upper = np.array(config["upper"], dtype=np.uint8)
            mask = cv2.inRange(hsv, lower, upper)

        # 形态学操作去噪
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (13, 13))
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

        cv2.imshow(color_name,mask)

        # 查找轮廓
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        for cnt in contours:
            area = cv2.contourArea(cnt)
            if  100 < area < 200:  # 过滤小面积噪声
                continue

            # 多边形近似
            peri = cv2.arcLength(cnt, True)
            epsilon = 0.01 * peri
            approx = cv2.approxPolyDP(cnt, epsilon, True)
            num_vertices = len(approx)

            if  8 >= num_vertices >= 4:  # 筛选3~8边形
                # 计算中心点
                M = cv2.moments(cnt)
                if M["m00"] != 0:
                    cx = int(M["m10"] / M["m00"])
                    cy = int(M["m01"] / M["m00"])
                else:
                    cx, cy = approx[0][0]

                # 绘制轮廓和中心点
                cv2.drawContours(result, [approx], -1, (0, 255, 0), 2)
                cv2.circle(result, (cx, cy), 5, (255, 0, 0), -1)

                # 标注颜色和顶点数
                text = f"{color_name} ({num_vertices})"
                cv2.putText(result, text, (cx - 20, cy - 20),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

    return result





capture = cv2.VideoCapture(1)  # 打开笔记本内置摄像头
while (capture.isOpened()):  # 笔记本内置摄像头被打开后
    retval, image = capture.read(1)  # 从摄像头中实时读取视频
    # img用做绘图，image用作处理
    img = image.copy()


    result_image = detect_colored_polygons(image)

    cv2.imshow("Video", result_image)  # 在窗口中显示读取到的视频

    key = cv2.waitKey(1)  # 窗口的图像刷新时间为1毫秒
    if key == 27:  # 如果按下esc
        break
capture.release()  # 关闭笔记本内置摄像头
cv2.destroyAllWindows()  # 销毁显示摄像头视频的窗口

